// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: yuxuan

#pragma once

#include <cmath>
#include <deque>
#include <utility>

#include "cc/coin2/strategy/util/moving_window.h"

namespace coin2::strategy::util {

template <typename T>
class MovingLm : public MovingWindow<T, std::pair<int64_t, int64_t>> {
 public:
  MovingLm(
      int64_t timewindow,
      int64_t start_ts,
      int64_t multiplier = 100,
      int64_t tmultiplier = 1e8,
      int64_t sample_interval = 5e8)
      : MovingWindow<T, std::pair<int64_t, int64_t>>(timewindow),
        start_ts_(start_ts),
        multiplier_(multiplier),
        tmultiplier_(tmultiplier),
        sample_interval_(sample_interval) {}

  void Add(int64_t ts, const T& y) final {
    Update(ts);
    int64_t iy = int64_t(y * multiplier_ + 0.4999);
    int64_t ix = TransformTs_(ts);
    y_ += iy;
    x_ += ix;
    x2_ += ix * ix;
    xy_ += iy * ix;
    MovingWindow<T, std::pair<int64_t, int64_t>>::dq_.push_back({ts, {ix, iy}});
    if (ts - last_sample_ts_ >= sample_interval_) {
      samples_.push_back({ts, 1.0 * iy / multiplier_});
      last_sample_ts_ = ts;
      if (samples_.size() > 1) {
        UpdateMse();
      }
    }
  }

  void Update(int64_t ts) final {
    while (!MovingWindow<T, std::pair<int64_t, int64_t>>::Empty() &&
           ts - MovingWindow<T, std::pair<int64_t, int64_t>>::dq_.front().first >
               MovingWindow<T, std::pair<int64_t, int64_t>>::timewindow_) {
      const auto& e = MovingWindow<T, std::pair<int64_t, int64_t>>::GetFirstData();
      int64_t ix = e.first;
      int64_t iy = e.second;
      y_ -= iy;
      x_ -= ix;
      x2_ -= ix * ix;
      xy_ -= ix * iy;
      MovingWindow<T, std::pair<int64_t, int64_t>>::dq_.pop_front();
    }
    while (!samples_.empty() && ts - samples_.front().first >
                                    MovingWindow<T, std::pair<int64_t, int64_t>>::timewindow_) {
      samples_.pop_front();
    }
  }

  void UpdateMse() {
    double sum = 0;
    double sum2 = 0;
    double diff = 0;
    for (const auto& p : samples_) {
      diff = Predict(p.first) - p.second;
      sum += diff;
      sum2 += diff * diff;
    }
    int n = samples_.size();
    sum /= n;
    double std = std::sqrt((sum2 - n * sum * sum) / (n - 1));
    mse_ = std;
  }

  double Predict(int64_t ts) {
    int n = MovingWindow<T, std::pair<int64_t, int64_t>>::Size();
    b_ = 1.0 * (n * xy_ - x_ * y_) / (n * x2_ - x_ * x_);
    a_ = 1.0 * (y_ - b_ * x_) / n;
    return (b_ * TransformTs_(ts) + a_) / multiplier_;
  }

  double Mse() { return mse_; }

 private:
  int64_t TransformTs_(int64_t ts) { return int64_t((ts - start_ts_ + 0.4999) / tmultiplier_); }

  int64_t start_ts_;
  int64_t multiplier_;
  int64_t tmultiplier_;
  int64_t sample_interval_;
  int64_t last_sample_ts_{0};
  std::deque<std::pair<int64_t, double>> samples_;
  int64_t y_{0};
  int64_t x_{0};
  int64_t x2_{0};
  int64_t xy_{0};
  double a_;
  double b_;
  double mse_;
};

}  // namespace coin2::strategy::util
